use ndarray::prelude::*;
use std::f64::INFINITY;

#[derive(Debug, Clone)]
pub struct DoseStatRet {
    pub binned_dose_values: Vec<f64>,
    pub accumulated_relative_volumes: Vec<f64>,
}

pub fn accumulate(voxel_volume: &[f64], voxel_dose: &[f64], do_exact: bool) -> DoseStatRet {
    let n = voxel_volume.len();
    let mut pairs: Vec<(f64, f64)> = voxel_volume
        .iter()
        .zip(voxel_dose.iter())
        .filter(|(vol, _)| **vol > 0.0)
        .map(|(vol, dose)| (dose.max(0.0), *vol))
        .collect();

    if pairs.is_empty() {
        return DoseStatRet {
            binned_dose_values: vec![0.0, 0.0],
            accumulated_relative_volumes: vec![1.0, 0.0],
        };
    }

    const EXACT_THRESHOLD: usize = 5000;
    let use_exact = do_exact || pairs.len() < EXACT_THRESHOLD;

    if use_exact {
        pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
        let mut doses = vec![0.0];
        let mut vols = vec![0.0];
        for (d, v) in pairs {
            doses.push(d);
            vols.push(v);
        }
        let total_vol: f64 = vols.iter().sum();
        let mut acc = vec![0.0; vols.len()];
        let mut cum = 0.0;
        for i in (1..vols.len()).rev() {
            cum += vols[i];
            acc[i] = cum;
        }
        for v in &mut acc {
            *v /= cum;
        }
        return DoseStatRet {
            binned_dose_values: doses,
            accumulated_relative_volumes: acc,
        };
    }

    // Rough binning
    let total_vol: f64 = pairs.iter().map(|(_, v)| *v).sum();
    let doses: Vec<f64> = pairs.iter().map(|(d, _)| *d).collect();
    let min_dose = doses.iter().fold(INFINITY, |a, &b| a.min(b)).max(0.0);
    let max_dose = *doses
        .iter()
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();
    let dose_range = max_dose - min_dose;
    const MIN_WIDTH: f64 = 1e-6;
    if dose_range <= MIN_WIDTH {
        panic!("Dose range too small for rough binning.");
    }

    let num_bins = ((dose_range / MIN_WIDTH) as usize).min(1001).max(2);
    let bin_width = dose_range / (num_bins - 1) as f64;

    let mut binned_doses = vec![0.0; num_bins + 2];
    let mut binned_vols = vec![0.0; num_bins + 2];

    for i in 0..num_bins {
        binned_doses[i + 1] = min_dose + i as f64 * bin_width;
    }

    for (d, v) in pairs {
        if d <= min_dose {
            binned_vols[1] += v;
        } else if d >= max_dose {
            // skip sentinel
        } else {
            let idx = ((d - min_dose) / bin_width + 0.5) as usize + 1;
            binned_vols[idx.min(num_bins - 1) + 1] += v;
        }
    }

    let mut cum = 0.0;
    for i in (1..binned_vols.len() - 1).rev() {
        let vol = binned_vols[i];
        binned_vols[i] = cum + 0.5 * vol;
        cum += vol;
    }
    binned_vols[0] = cum;
    *binned_vols.last_mut().unwrap() = 0.0;

    for v in &mut binned_vols {
        *v /= cum;
    }

    binned_doses[0] = binned_doses[1];
    *binned_doses.last_mut().unwrap() = binned_doses[num_bins];

    DoseStatRet {
        binned_dose_values: binned_doses,
        accumulated_relative_volumes: binned_vols,
    }
}

pub fn get_metric(vol_pct: f64, dose_values: &[f64], acc_vols: &[f64]) -> f64 {
    if dose_values.is_empty() || acc_vols.is_empty() {
        return 0.0;
    }
    if vol_pct >= acc_vols[0] {
        return dose_values[0];
    }
    if vol_pct <= *acc_vols.last().unwrap() {
        return *dose_values.last().unwrap();
    }

    let pos = acc_vols.binary_search_by(|&v| v.partial_cmp(&vol_pct).unwrap().reverse());
    match pos {
        Ok(i) => dose_values[i],
        Err(i) => {
            let i1 = (i > 0).then(|| i - 1).unwrap_or(0);
            let i2 = (i < acc_vols.len())
                .then(|| i)
                .unwrap_or(acc_vols.len() - 1);
            let v1 = acc_vols[i1];
            let v2 = acc_vols[i2];
            let d1 = dose_values[i1];
            let d2 = dose_values[i2];
            if (v2 - v1).abs() < 1e-10 {
                (d1 + d2) / 2.0
            } else {
                d1 + (vol_pct - v1) * (d2 - d1) / (v2 - v1)
            }
        }
    }
}
